
#ifndef ROPUFU_MULTI_SENSOR_MONTE_CARLO_LOOP_HPP_INCLUDED
#define ROPUFU_MULTI_SENSOR_MONTE_CARLO_LOOP_HPP_INCLUDED

#include <aftermath/probability/empirical_measure.hpp>

#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>

namespace ropufu
{
    namespace monte_carlo
    {
        namespace detail
        {
            template <typename t_statistic_type, typename t_stopping_time_type>
            struct loop
            {
                typedef t_statistic_type statistic_type;
                //typedef typename statistic_type::information_type info_type;
                //typedef typename statistic_type::window_type window_type;

                typedef t_stopping_time_type stopping_time_type;
                typedef loop<statistic_type, stopping_time_type> type;

                static const std::size_t count_streams = statistic_type::count_streams;

                typedef aftermath::probability::empirical_measure<stopping_time_type, double> empirical_time_type;
                typedef aftermath::probability::empirical_measure<double, double> empirical_double_type;
                typedef aftermath::probability::empirical_measure<bool> empirical_bool_type;

            private:
                std::uint_fast64_t m_number_of_runs;
                std::size_t m_count_affected;

                std::vector<double> m_threshold;

                // For each threshold we have an empirical measure for stopping time.
                std::vector<empirical_time_type> m_run_length;
                // For each threshold we have an auxiliary scalar observation on stopping.
                std::vector<empirical_double_type> m_aux;
                // For each threshold we have a flag if the procedure stopped because of a reset.
                std::vector<empirical_double_type> m_reset_flag;
                // For each threshold store a flag whether the respective procedure is still running.
                std::vector<bool> m_is_running;

            public:
                explicit loop(std::uint_fast64_t number_of_runs, std::size_t count_affected, const std::vector<double>& threshold)
                    : m_number_of_runs(number_of_runs), m_count_affected(count_affected), m_threshold(threshold),
                    m_run_length(threshold.size()), m_aux(threshold.size()), m_reset_flag(threshold.size()), m_is_running(threshold.size())
                {

                }

                template <typename t_uniform_engine_type, typename t_sampler_type>
                void execute(statistic_type& procedure, bool do_stop_on_reset, 
                    t_uniform_engine_type& uniform_engine, t_sampler_type& sampler_before, t_sampler_type& sampler_after)
                {
                    // For each threshold store a flag whether the respective procedure is still running.
                    for (std::uint_fast64_t i = 0; i < this->m_number_of_runs; i++)
                    {
                        stopping_time_type time = 0;
                        procedure.reset();

                        // Initialize flags with 1-s.
                        this->m_is_running = std::vector<bool>(this->m_threshold.size(), true);
                        // Keep track of how many thresholds are to be crossed.
                        std::size_t count_alive = this->m_threshold.size();

                        double current_stat = 0.0;
                        while (count_alive > 0)
                        {
                            // l_t, the current instantaneous log-lieklihood ratio.
                            std::array<double, count_streams> current_loglikelihood;
                            for (std::size_t k = 0; k < this->m_count_affected; k++) current_loglikelihood[k] = sampler_after(uniform_engine);
                            for (std::size_t k = this->m_count_affected; k < count_streams; k++) current_loglikelihood[k] = sampler_before(uniform_engine);

                            current_stat = procedure.observe(current_loglikelihood);
                            time++;

                            // Check for crossed thresholds.
                            for (std::size_t j = 0; j < this->m_threshold.size(); j++)
                            {
                                double a = this->m_threshold[j];
                                bool has_crossed = current_stat >= a;
                                bool has_reset = procedure.has_reset();
                                bool do_stop = has_crossed || (has_reset && do_stop_on_reset);
                                if (this->m_is_running[j] && do_stop)
                                {
                                    // Record the stopping time for current threshold, a.
                                    this->m_run_length[j].observe(time);
                                    if (do_stop_on_reset)
                                    {
                                        // Record the auxiliary function on stopping.
                                        const std::array<double, count_streams>& complete_loglikelihood = procedure.complete();
                                        double total_lr = 0.0;
                                        for (std::size_t k = 0; k < count_streams; k++) total_lr += std::exp(complete_loglikelihood[k]);
                                        
                                        this->m_aux[j].observe(has_reset ? 0.0 : (1.0 / total_lr));
                                        this->m_reset_flag[j].observe(has_reset ? 1.0 : 0.0);
                                    }
                                    // Make sure that we don't do it twice.
                                    this->m_is_running[j] = false;
                                    count_alive--;
                                }
                            }
                        }
                    }
                }

                const std::vector<empirical_time_type>& run_length() const
                {
                    return this->m_run_length;
                }

                const std::vector<empirical_double_type>& aux() const
                {
                    return this->m_aux;
                }

                const std::vector<empirical_double_type>& reset_flag() const
                {
                    return this->m_reset_flag;
                }
            };
        }
    }
}

#endif // ROPUFU_MULTI_SENSOR_MONTE_CARLO_LOOP_HPP_INCLUDED
